from typing import List, Dict, Any, Optional, Union, Literal
from pydantic import BaseModel, Field, validator
# import numpy as np # 对于 Pydantic 模型定义并非严格需要，但有助于示例

# --- 辅助模型 ---

class MeshParameters(BaseModel):
    """网格生成参数。"""
    mesh_type: str = Field(default="box", description="网格类型 (例如, 'box', 'gmsh_file')")
    nx: Optional[int] = Field(default=10, description="box 网格在 x 方向上的单元数量")
    ny: Optional[int] = Field(default=10, description="box 网格在 y 方向上的单元数量")
    nz: Optional[int] = Field(default=5, description="box 网格在 z 方向上的单元数量")
    lx: Optional[float] = Field(default=0.1, description="box 网格在 x 方向上的长度 [m]")
    ly: Optional[float] = Field(default=0.1, description="box 网格在 y 方向上的长度 [m]")
    lz: Optional[float] = Field(default=0.001, description="box 网格在 z 方向上的长度 (厚度) [m]")
    gmsh_file_path: Optional[str] = Field(default=None, description="如果 mesh_type 为 'gmsh_file', 则为 GMSH .msh 文件的路径")
    scale_factor: float = Field(default=1.0, description="网格坐标的缩放因子")

class DomainConfig(BaseModel):
    """网格内单个域的配置。"""
    name: str = Field(..., description="域的名称 (例如, 'anode_gdl', 'cathode_cl', 'membrane')")
    id: int = Field(..., description="在网格中标记此域的整数 ID")
    material: str = Field(..., description="分配给此域的材料名称 (链接到 MaterialProperties)")

class BoundaryConditionConfig(BaseModel):
    """边界条件的配置。"""
    name: str = Field(..., description="边界的名称 (例如, 'anode_current_collector', 'cathode_outlet')")
    marker_id: int = Field(..., description="在 mesh_tags 中标记此边界的整数 ID")
    bc_type: str = Field(..., description="边界条件的类型 (例如, 'dirichlet', 'neumann', 'robin')")
    variable: str = Field(..., description="此边界条件适用的变量名称 (例如, 'phi_s', 'phi_m', 'T', 'C_O2')")
    value: Union[float, List[float], str] = Field(..., description="Dirichlet 边界条件的值。对于 Neumann/Robin，其解释取决于具体实现。")
    expression_coeffs: Optional[Dict[str, float]] = Field(default=None, description="如果 value 是字符串表达式，则为基于表达式的边界条件的系数。")

class OperatingConditions(BaseModel):
    """模拟的操作条件。"""
    cell_voltage: float = Field(default=0.7, description="电池电压 [V] (用作阴极处 phi_s 的 Dirichlet 条件)")
    temperature_inlet: float = Field(default=343.15, description="入口/边界温度 [K]")
    pressure_anode_inlet: float = Field(default=1.5e5, description="阳极入口压力 [Pa]")
    pressure_cathode_inlet: float = Field(default=1.5e5, description="阴极入口压力 [Pa]")
    anode_stoichiometry: float = Field(default=1.5, description="H2 的阳极化学计量数")
    cathode_stoichiometry: float = Field(default=2.0, description="O2/空气 的阴极化学计量数")
    relative_humidity_anode: float = Field(default=1.0, description="阳极入口的相对湿度 (0 到 1)")
    relative_humidity_cathode: float = Field(default=1.0, description="阴极入口的相对湿度 (0 到 1)")

class MaterialParameters(BaseModel):
    """材料的物理和电化学参数。"""
    # 通用属性
    density: Optional[float] = Field(default=None, description="密度 [kg/m^3]")
    heat_capacity: Optional[float] = Field(default=None, description="比热容 [J/(kg*K)]")
    thermal_conductivity_eff: Optional[float] = Field(default=None, description="有效导热系数 [W/(m*K)]")
    
    # 电学属性
    solid_phase_conductivity_eff: Optional[float] = Field(default=None, description="有效固相 (电子) 电导率 [S/m]")
    membrane_phase_conductivity_eff: Optional[float] = Field(default=None, description="有效膜相 (质子) 电导率 [S/m]") # 通常取决于含水量
    
    # 电化学属性 (通常用于催化层 CLs)
    exchange_current_density_HOR_ref: Optional[float] = Field(default=None, description="HOR 的参考交换电流密度 [A/m^2_Pt]")
    exchange_current_density_ORR_ref: Optional[float] = Field(default=None, description="ORR 的参考交换电流密度 [A/m^2_Pt]")
    alpha_anodic_HOR: Optional[float] = Field(default=None, description="HOR 的阳极传递系数")
    alpha_cathodic_HOR: Optional[float] = Field(default=None, description="HOR 的阴极传递系数")
    alpha_anodic_ORR: Optional[float] = Field(default=None, description="ORR 的阳极传递系数")
    alpha_cathodic_ORR: Optional[float] = Field(default=None, description="ORR 的阴极传递系数")
    reference_concentration_H2_ref: Optional[float] = Field(default=None, description="HOR 的参考 H2 浓度 [mol/m^3]")
    reference_concentration_O2_ref: Optional[float] = Field(default=None, description="ORR 的参考 O2 浓度 [mol/m^3]")
    active_surface_area_ratio: Optional[float] = Field(default=None, description="单位体积的活性表面积 (Av) [m^2_Pt/m^3_CL]")
    pt_loading_mg_cm2: Optional[float] = Field(default=None, description="铂载量 [mg_Pt/cm^2_electrode]")
    
    # 气体传输属性 (多孔介质)
    porosity: Optional[float] = Field(default=None, description="孔隙率 (epsilon)")
    tortuosity: Optional[float] = Field(default=None, description="曲折度 (tau)")
    permeability: Optional[float] = Field(default=None, description="渗透率 (Darcy) [m^2]")
    
    # 水传输属性 (膜/离聚物)
    electro_osmotic_drag_coeff: Optional[float] = Field(default=None, description="电渗拖曳系数 (xi)")
    water_diffusivity_membrane_ref: Optional[float] = Field(default=None, description="膜中参考水扩散系数 [m^2/s]")
    
    # 其他模型特定参数
    model_specific_params: Dict[str, Any] = Field(default_factory=dict, description="其他模型特定参数")

class MaterialProperties(BaseModel):
    """所有材料定义的容器。"""
    materials: Dict[str, MaterialParameters] = Field(default_factory=dict, description="材料名称到其参数的字典")

class CouplingSettings(BaseModel):
    """物理场耦合求解的设置"""
    max_iterations: int = Field(100, description="物理场耦合求解的最大迭代次数")
    absolute_tolerance: float = Field(1.0e-6, description="物理场耦合求解的绝对收敛容差")
    relative_tolerance: float = Field(1.0e-4, description="物理场耦合求解的相对收敛容差")
    relaxation_factor: float = Field(0.7, description="欠松弛因子 (0 < factor <= 1)")
    monitor_residual: bool = Field(True, description="是否监视并打印每个耦合迭代的残差")
    save_intermediate_results: bool = Field(False, description="是否保存中间耦合步骤的结果")
    save_frequency: int = Field(10, description="如果 save_intermediate_results 为 True，则指定保存结果的频率 (每 N 次迭代)")
    coupling_strategy: Literal["picard", "newton_coupled"] = Field("picard", description="多物理场耦合策略")

class LinearSolverSettings(BaseModel):
    """线性求解器的特定设置"""
    solver_type: Literal["ksp", "gmres", "cg", "bicgstab", "minres"] = Field("gmres", description="线性求解器类型 (PETSc KSP 类型)")
    preconditioner_type: Literal["ilu", "jacobi", "sor", "asm", "gamg", "hypre", "none"] = Field("ilu", description="预处理器类型 (PETSc PC 类型)")
    max_iterations: int = Field(1000, description="线性求解器的最大迭代次数")
    absolute_tolerance: float = Field(1.0e-9, description="线性求解器的绝对收敛容差")
    relative_tolerance: float = Field(1.0e-6, description="线性求解器的相对收敛容差")
    divergence_tolerance: float = Field(1.0e4, description="线性求解器的发散容差")
    monitor_convergence: bool = Field(False, description="是否监视线性求解器的收敛过程 (例如 -ksp_monitor)")
    initial_guess_nonzero: bool = Field(True, description="是否使用非零初始猜测 (如果可用)")
    custom_options: Optional[Dict[str, Union[str, float, int, bool]]] = Field(None, description="PETSc 选项的自定义字典，例如 {'ksp_gmres_restart': 30}")

class NonlinearSolverSettings(BaseModel):
    """非线性求解器的特定设置"""
    solver_type: Literal["newton", "snes"] = Field("newton", description="非线性求解器类型 (FEniCSx NewtonSolver 或 PETSc SNES)")
    max_iterations: int = Field(50, description="非线性求解器的最大迭代次数")
    absolute_tolerance: float = Field(1.0e-9, description="非线性求解器的绝对收敛容差")
    relative_tolerance: float = Field(1.0e-6, description="非线性求解器的相对收敛容差")
    relaxation_parameter: float = Field(1.0, description="牛顿法的松弛参数 (阻尼因子)")
    report: bool = Field(True, description="打印牛顿求解器的收敛报告 (FEniCSx NewtonSolver)")
    error_on_nonconvergence: bool = Field(False, description="如果牛顿求解器未收敛则抛出错误 (FEniCSx NewtonSolver)")
    snes_type: Literal["newtonls", "newtontr", "richardson", "ngmres", "fas"] = Field("newtonls", description="PETSc SNES 类型 (如果 solver_type 为 'snes')")
    monitor_convergence: bool = Field(False, description="是否监视 SNES 求解器的收敛过程 (例如 -snes_monitor)")
    custom_options: Optional[Dict[str, Union[str, float, int, bool]]] = Field(None, description="PETSc 选项的自定义字典，例如 {'snes_linesearch_type': 'basic'}")

class SolverSettings(BaseModel):
    """单个物理问题求解器的设置"""
    solver_type: Literal["linear", "nonlinear"] = Field("nonlinear", description="求解器类型: 'linear' 或 'nonlinear'")
    # 以下参数是通用的，但主要影响非线性求解器或耦合迭代的全局行为。对于纯线性求解，部分可能不直接使用。
    max_iterations: int = Field(50, description="求解器最大迭代次数 (适用于牛顿法或类似方法)")
    absolute_tolerance: float = Field(1.0e-9, description="绝对收敛容差")
    relative_tolerance: float = Field(1.0e-6, description="相对收敛容差")
    relaxation_factor: float = Field(1.0, description="松弛因子 (如果适用，例如用于牛顿法的阻尼)") # 注意：NonlinearSolverSettings 中也有 relaxation_parameter
    report_details: bool = Field(False, description="是否报告详细的求解器统计信息")

    linear_solver_config: Optional[LinearSolverSettings] = Field(None, description="线性求解器的具体配置 (如果 solver_type 为 'linear' 或非线性求解器内部使用线性求解)")
    nonlinear_solver_config: Optional[NonlinearSolverSettings] = Field(None, description="非线性求解器的具体配置 (如果 solver_type 为 'nonlinear')")

    @validator('linear_solver_config', pre=True, always=True)
    def set_linear_solver_config(cls, v, values):
        # 如果 solver_type 是 'linear'，确保 linear_solver_config 被实例化
        # 或者如果 nonlinear_solver_config 存在 (暗示可能需要内部线性求解)
        solver_type = values.get('solver_type')
        if solver_type == 'linear' and v is None:
            return LinearSolverSettings()
        # 如果是非线性求解器，它内部也可能需要线性求解器设置，所以如果没提供，也给个默认的
        # 更好的做法可能是在 NonlinearSolverWrapper 中处理其内部线性求解器的配置
        # 但为了简单起见，这里如果是非线性且没提供，也给一个
        if solver_type == 'nonlinear' and v is None: 
            # 实际上，非线性求解器通常会配置其内部的线性求解器。这里可能不需要自动创建。
            # 暂时注释掉，让用户在需要时显式提供。
            # return LinearSolverSettings() 
            pass # 让它保持 None，除非显式提供
        return v

    @validator('nonlinear_solver_config', pre=True, always=True)
    def set_nonlinear_solver_config(cls, v, values):
        if values.get('solver_type') == 'nonlinear' and v is None:
            return NonlinearSolverSettings()
        return v

    class Config:
        validate_assignment = True # 当字段被重新赋值时进行验证

class TimeSteppingSettings(BaseModel):
    """时间依赖性模拟的设置。"""
    total_time: float = Field(default=1.0, description="总模拟时间 [s]")
    num_time_steps: Optional[int] = Field(default=None, description="时间步数 (必须设置此项或 dt)")
    dt_initial: Optional[float] = Field(default=None, description="初始时间步长 [s] (必须设置此项或 num_time_steps)")
    dt_min: Optional[float] = Field(default=1e-6, description="最小自适应时间步长 [s]")
    dt_max: Optional[float] = Field(default=0.1, description="最大自适应时间步长 [s]")
    adaptive_stepping: bool = Field(default=False, description="启用自适应时间步进")
    time_integration_scheme: str = Field(default="BE", description="时间积分方案 (例如, 'BE' 代表后向欧拉, 'CN' 代表 Crank-Nicolson)")

class OutputSettings(BaseModel):
    """模拟输出的设置。"""
    output_path: str = Field(default="results/", description="保存输出文件的目录")
    output_frequency_steps: Optional[int] = Field(default=10, description="每 N 个求解器步骤/迭代输出结果")
    output_frequency_time: Optional[float] = Field(default=None, description="每 T 秒模拟时间输出结果")
    save_vtx_files: bool = Field(default=True, description="以 VTX/BP 格式保存解，用于 ParaView/VisIt")
    save_xdmf_files: bool = Field(default=False, description="以 XDMF 格式保存解")
    save_csv_data: bool = Field(default=False, description="将标量数据 (例如, 电流, 平均值) 保存到 CSV")
    pyvista_visualization: bool = Field(default=False, description="结束时启用 PyVista 可视化 (需要桌面环境)")
    
class Constants(BaseModel):
    """物理和数学常数。"""
    faraday_constant_C_per_mol: float = Field(default=96485.33212, description="法拉第常数 (F) [C/mol]")
    gas_constant_J_per_mol_K: float = Field(default=8.314462618, description="通用气体常数 (R) [J/(mol*K)]")
    standard_temperature_K: float = Field(default=298.15, description="标准温度 [K]")
    standard_pressure_Pa: float = Field(default=101325.0, description="标准压力 [Pa]")

# --- 添加 NumericalSettings 定义 ---
class NumericalSettings(BaseModel):
    """控制数值方法和参数的设置。"""
    # 有限元相关
    default_fem_degree: int = Field(1, description="默认的有限元函数空间多项式阶数")
    default_quadrature_degree: Optional[int] = Field(None, description="用于积分的默认求积阶数 (如果为None，则自动确定)")
    # 几何与网格处理
    geometric_tolerance: float = Field(1e-6, description="几何操作的容差 (例如，点定位)")
    # 其他可能的数值调整参数
    stabilization_factor: Optional[float] = Field(None, description="如果使用数值稳定化方法 (如SUPG)，则为其参数")
# --- 结束 NumericalSettings 定义 ---

# --- 主要 SimulationParameters 模型 ---

class SimulationParameters(BaseModel):
    """所有模拟参数的主要 Pydantic 模型。"""
    project_name: str = Field(default="PEMFC Simulation", description="模拟项目的名称")
    description: str = Field(default="A FEniCSx-based PEM Fuel Cell simulation", description="项目的简要描述")
    version: str = Field(default="0.1.0", description="项目版本")

    mesh_parameters: MeshParameters = Field(default_factory=MeshParameters, description="网格生成参数")
    domains: List[DomainConfig] = Field(default_factory=list, description="域配置列表")
    boundary_conditions: List[BoundaryConditionConfig] = Field(default_factory=list, description="边界条件配置列表")
    
    operating_conditions: OperatingConditions = Field(default_factory=OperatingConditions, description="电池操作条件")
    material_properties: MaterialProperties = Field(default_factory=MaterialProperties, description="材料属性定义")
    
    solver_settings: SolverSettings = Field(default_factory=SolverSettings, description="数值求解器设置")
    time_stepping: Optional[TimeSteppingSettings] = Field(default=None, description="时间步进设置 (用于瞬态模拟)")
    output_settings: OutputSettings = Field(default_factory=OutputSettings, description="输出和可视化设置")
    
    physical_constants: Constants = Field(default_factory=Constants, description="要使用的物理常数")
    coupling_settings: Optional[CouplingSettings] = Field(default_factory=CouplingSettings, description="多物理场耦合设置")
    numerical_settings: NumericalSettings = Field(default_factory=NumericalSettings, description="数值方法和参数")

    # 模型特定的标志或选择
    enable_thermal_model: bool = Field(default=True, description="启用或禁用热模型 (能量平衡) 的标志")
    enable_gas_transport: bool = Field(default=False, description="启用或禁用气相组分传输的标志")
    enable_water_transport: bool = Field(default=False, description="启用或禁用液态水传输模型 (在膜/离聚物中) 的标志")
    
    # 有限元设置
    fem_degree_phi_s: int = Field(default=1, description="固相电势 (phi_s) 的有限元阶数")
    fem_degree_phi_m: int = Field(default=1, description="膜相电势 (phi_m) 的有限元阶数")
    fem_degree_T: int = Field(default=1, description="温度 (T) 的有限元阶数")
    fem_degree_gas_concentration: int = Field(default=1, description="气体浓度的有限元阶数")
    fem_degree_water_content: int = Field(default=1, description="含水量/水活度的有限元阶数")
    
    # 其他全局设置
    verbosity_level: int = Field(default=1, ge=0, le=2, description="详细级别 (0: 静默, 1: 正常, 2: 调试)")
    
    model_config = {
        "validate_assignment": True, # 确保在初始化后分配值时进行验证
        "extra": "forbid" # 禁止模型中未定义的额外字段
    }

if __name__ == '__main__':
    # 示例用法:
    # 创建一个默认参数实例
    default_params = SimulationParameters()
    print("--- 默认参数 ---")
    print(default_params.model_dump_json(indent=2))

    # 创建一个自定义参数实例
    custom_params_data = {
        "project_name": "我的 PEMFC 测试",
        "mesh_parameters": {"nx": 20, "nz": 10, "lz": 0.0002},
        "domains": [
            {"name": "anode_gdl", "id": 1, "material": "CarbonPaper"},
            {"name": "anode_cl", "id": 2, "material": "AnodeCatalystLayer"},
            {"name": "membrane", "id": 3, "material": "Nafion117"},
            {"name": "cathode_cl", "id": 4, "material": "CathodeCatalystLayer"},
            {"name": "cathode_gdl", "id": 5, "material": "CarbonPaper"},
        ],
        "material_properties": {
            "materials": {
                "CarbonPaper": {"solid_phase_conductivity_eff": 5000.0, "thermal_conductivity_eff": 1.5, "porosity": 0.7},
                "AnodeCatalystLayer": {
                    "solid_phase_conductivity_eff": 1000.0, 
                    "membrane_phase_conductivity_eff": 0.5, # 简化
                    "thermal_conductivity_eff": 1.0,
                    "active_surface_area_ratio": 2e5,
                    "exchange_current_density_HOR_ref": 1e-1, # A/m^2_Pt (示例值)
                    "alpha_anodic_HOR": 0.5,
                    "alpha_cathodic_HOR": 0.5,
                    "reference_concentration_H2_ref": 10.0, # mol/m^3
                    "porosity": 0.4
                },
                 "Nafion117": {"membrane_phase_conductivity_eff": 1.0, "thermal_conductivity_eff": 0.3}, # 简化
                 "CathodeCatalystLayer": {
                    "solid_phase_conductivity_eff": 1000.0, 
                    "membrane_phase_conductivity_eff": 0.5, # 简化
                    "thermal_conductivity_eff": 1.0,
                    "active_surface_area_ratio": 5e5,
                    "exchange_current_density_ORR_ref": 1e-5, # A/m^2_Pt (示例值)
                    "alpha_cathodic_ORR": 0.5,
                    "reference_concentration_O2_ref": 4.0, # mol/m^3
                    "porosity": 0.4
                },
            }
        },
        "operating_conditions": {"cell_voltage": 0.65},
        "solver_settings": {"absolute_tolerance": 1e-7},
        "output_settings": {"output_path": "results_my_test", "pyvista_visualization": True}
    }
    custom_params = SimulationParameters(**custom_params_data)
    print("\n--- 自定义参数 ---")
    print(custom_params.model_dump_json(indent=2))

    # 访问嵌套参数的示例
    print(f"\n自定义电池电压: {custom_params.operating_conditions.cell_voltage}")
    print(f"阳极催化层材料: {custom_params.domains[1].material}")
    print(f"阳极催化层孔隙率: {custom_params.material_properties.materials['AnodeCatalystLayer'].porosity}")

    # YAML 加载器如何使用它的示例 (概念性)
    """
    import yaml
    with open("config/input_params.yaml", 'r') as f:
        config_dict = yaml.safe_load(f)
        loaded_params = SimulationParameters(**config_dict)
        print("
--- 从概念性 YAML 加载的参数 ---")
        print(loaded_params.model_dump_json(indent=2))
    """ 